class NumSequence:
    PAD_TAG = "PAD"  # 补的字符
    PAD = 0

    UNK_TAG = "UNK"  # 未知的字符
    UNK = 1

    SOS_TAG = "SOS"  # 起始字符
    SOS = 2

    EOS_TAG = "EOS"  # 结束字符
    EOS = 3

    def __init__(self):
        self.dict = {self.PAD_TAG: self.PAD, self.UNK_TAG: self.UNK, self.SOS_TAG: self.SOS, self.EOS_TAG: self.EOS}
        for i in range(10):
            self.dict[str(i)] = len(self.dict)
        self.inverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def transform(self, sentence, max_len=10, add_eos=False):
        if len(sentence) > max_len:
            sentence = sentence[:max_len]
        sentence_len = len(sentence)
        if add_eos:
            sentence = sentence + [self.EOS_TAG]
        if sentence_len < max_len:
            sentence = sentence + [self.PAD_TAG] * (max_len - sentence_len)
        result = [self.dict.get(i, self.UNK) for i in sentence]
        return result

    def inverse_transform(self, indices):
        return [self.inverse_dict.get(i, self.UNK_TAG) for i in indices]

        # result = []
        # for i in indices:
        #     if i == self.EOS:
        #         break
        #     result.append(self.inverse_dict.get(int(i), self.UNK_TAG))
        # return result

    def __len__(self):
        return len(self.dict)
